# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import atexit
import concurrent.futures
import pickle
import gzip
import sys
import threading

import asyncio
import cloudpickle
import six
from tornado.ioloop import IOLoop, PeriodicCallback
from tornado.process import Subprocess as TornadoSubprocess
from tornado.platform.asyncio import AnyThreadEventLoopPolicy


def write_file(data, file, flush=True, compress=False, compresslevel=5):
  if compress:
    with gzip.GzipFile(mode='wb', compresslevel=compresslevel, fileobj=file) as gzip_file:
      gzip_file.write(data)
  else:
    file.write(data)

  if flush:
    file.flush()


def pickle_dump_to_file(obj, filepath, compresslevel=5):
  if filepath.endswith('.gz'):
    fileobj = gzip.open(filepath, 'wb', compresslevel=compresslevel)
  else:
    fileobj = open(filepath, 'wb')
  with fileobj:
    cloudpickle.dump(obj, fileobj)


def pickle_load_from_file(filepath):
  if filepath.endswith('.gz'):
    fileobj = gzip.open(filepath, 'rb')
  else:
    fileobj = open(filepath, 'rb')
  with fileobj:
    return pickle.load(fileobj)


class Future(concurrent.futures.Future):
  def __init__(self):
    super().__init__()
    self._exc_info = None

  def set_exception_info(self, exc_info=None):
    exc_info = exc_info or sys.exc_info()
    self._exc_info = exc_info
    self.set_exception(exc_info[1])

  def result(self):
    try:
      return super().result()
    except Exception:
      if self._exc_info:
        six.reraise(*self._exc_info)
      else:
        raise

  def exception_info(self):
    self.exception()
    return self._exc_info


class TornadoSubprocessWatcher(object):
  """Helps to use tornado.process.Subprocess without running IOLoop on the main
  thread. If user wants to use tornado subprocess in another thread, one needs
  write the below code in the main thread.

    TornadoSubprocessWatcher.install()

  Tornado's implementation to watch a child process uses SIGCHLD but due to the
  Python's bug (https://bugs.python.org/issue21895), main thread does no
  receive SIGCHLD signal from a child process executed in another thread.
  So here we make watcher to periodically watch its children rather than using
  signal.
  """
  _instance = None

  @classmethod
  def install(cls):
    if cls._instance is not None:
      return
    cls.instance = cls()

  def __init__(self):
    # This method modifies internal state of Tornado's Subprocess implementation.
    # https://www.tornadoweb.org/en/stable/_modules/tornado/process.html#Subprocess
    assert not TornadoSubprocess._initialized
    self._ioloop = None

    # Without this, IOLoop.current() should fail except the main thread
    asyncio.set_event_loop_policy(AnyThreadEventLoopPolicy())

    ready = threading.Event()
    self._thread = threading.Thread(target=self._thread_main, args=(ready,), daemon=True)
    self._thread.start()
    ready.wait()

    TornadoSubprocess._initialized = True
    atexit.register(self.stop)

  def stop(self):
    if self._thread is None:
      return
    self._ioloop.add_callback(self._ioloop.stop)
    self._thread.join()
    self._thread = None

  def _thread_main(self, ready):
    self._ioloop = IOLoop.current()
    ready.set()

    # Watch children every 100ms.
    PeriodicCallback(TornadoSubprocess._cleanup, 100).start()
    self._ioloop.start()
